from src.common.logger import getLogger
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

logger = getLogger()

class MultiQuery:

    def __init__(self, llm_model, vector_store):
        self.llm_model = llm_model
        self.vector_store = vector_store

    def invoke(self, query):
        logger.info(f"MultiQuery invoke query: {query}")
        retriever = self.vector_store.as_retriever(search_kwargs = { "k": 3 })

        multi_query_template = """
            你是一个智能助手AI，你的任务是用给定的问题生成3个不同版本的问题，使生成的3个问题的检索相似度与给定的问题一致。生成的3个问题用换行符隔开。
            给定的问题： {question}
        """
        multi_query_prompt = ChatPromptTemplate.from_template(multi_query_template)
        multi_query_chain = multi_query_prompt | self.llm_model | StrOutputParser() | (lambda x: x.split("\n"))
        multi_queries = multi_query_chain.invoke({ "question": query })
        logger.info(f"MultiQuery invoke multi_queries: {multi_queries}")

        multi_query_retrieve_chain = multi_query_chain | retriever.map() | self.merge_unique_document
        multi_query_retrieve_docs = multi_query_retrieve_chain.invoke({ "question": query })
        logger.info(f"MultiQuery invoke multi_query_retrieve_docs len: {len(multi_query_retrieve_docs)}")

        template = """
            请基于以下上下文内容回答问题：
            {context}
            问题：{question}
        """
        prompt = ChatPromptTemplate.from_template(template)
        chain = prompt | self.llm_model | StrOutputParser()
        chain_result = chain.invoke({ "context": multi_query_retrieve_docs, "question": query })
        logger.info(f"MultiQuery invoke chain_result len: {len(chain_result)}")
        return { "retrieve_docs": multi_query_retrieve_docs, "chain_result": chain_result }

    def merge_unique_document(self, documents):
        flattern_docs = [doc.page_content for subList in documents for doc in subList]
        uniques_docs = list(set(flattern_docs))
        return [doc for doc in uniques_docs]
